from langchain_community.vectorstores import Milvus
from langchain_core.embeddings import Embeddings
from langchain_core.documents import Document
from typing import List, Dict, Any, Optional
from app.config.settings import settings

class MilvusService:
    """Milvus向量数据库服务"""
    
    def __init__(self, embedding_model: Embeddings):
        """初始化Milvus服务
        
        Args:
            embedding_model: 用于文本嵌入的模型
        """
        self.embedding_model = embedding_model
        self.connection_args = {
            "host": settings.MILVUS_HOST,
            "port": settings.MILVUS_PORT,
        }
        self.collection_name = settings.MILVUS_COLLECTION
        self._vector_store = None
    
    @property
    def vector_store(self) -> Milvus:
        """获取或创建Milvus向量存储"""
        if self._vector_store is None:
            self._vector_store = Milvus(
                embedding_function=self.embedding_model,
                collection_name=self.collection_name,
                connection_args=self.connection_args,
            )
        return self._vector_store
    
    def add_documents(self, documents: List[Document]) -> None:
        """添加文档到向量存储
        
        Args:
            documents: 要添加的文档列表
        """
        self.vector_store.add_documents(documents)
    
    def similarity_search(self, query: str, k: int = 4) -> List[Document]:
        """执行相似性搜索
        
        Args:
            query: 查询文本
            k: 返回的最相似文档数量
            
        Returns:
            相似文档列表
        """
        return self.vector_store.similarity_search(query, k=k)
    
    def similarity_search_with_score(self, query: str, k: int = 4) -> List[tuple[Document, float]]:
        """执行带分数的相似性搜索
        
        Args:
            query: 查询文本
            k: 返回的最相似文档数量
            
        Returns:
            (文档, 相似度分数)元组的列表
        """
        return self.vector_store.similarity_search_with_score(query, k=k)